#include "../DWMRICore/DTI.h"
#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/Point3I.h"

#include <omp.h>

#include <string>

using namespace std;

char *pDTIPathname = NULL;				/* input diffusion tensor pathname */
char *pFiberPathname = NULL;			/* output fiber configure pathname */

float stepSize = 0.2f;
float maxLength = 50.0f;
float faThreshold = 0.08f;
float angleThreshold = 0.707f;
int scale = 1;
int innerSteps = 10;

CDTI *pDTI = NULL;

// constant
const float DELTA = 0.00001f;


void PrintUsage()
{
	printf("Usage: ");
	printf("TensorDeflection -dti <DTI> -fiber <Fiber> [-stepsize <stepsize> -maxlength <maxlength> -fa_thresold <fa_threshold> -angle_threshold <fa_threshold> -scale <scale>]\n");
	printf("\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-fiber") == 0) {
			pFiberPathname = argv[++i];
		} else if (strcmp(argv[i], "-dti") == 0) {
			pDTIPathname = argv[++i];
		} else if (strcmp(argv[i], "-stepsize") == 0) {
			stepSize = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-maxlength") == 0) {
			maxLength = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-fa_thresold") == 0) {
			faThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-angle_threshold") == 0) {
			angleThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-scale") == 0) { 
			scale = atoi(argv[++i]);
		} else {
			printf("unknown %s\n", argv[i]);
			return -1;
		}
	}

	if (pDTIPathname == NULL || pFiberPathname == NULL || maxLength < 0.0001f || stepSize < 0.00001f)
		return -2;

	return 0;
}


void ScaleDTI()
{
	CDTI *pTemp = new CDTI();
	pTemp->OpenDTIFile(pDTIPathname);

	uint w = pTemp->GetWidth();
	uint h = pTemp->GetHeight();
	uint d = pTemp->GetDepth();

	uint sw = (w - 1) * scale + 1;
	uint sh = (h - 1) * scale + 1;
	uint sd = (d - 1) * scale + 1;

	pDTI = new CDTI();
	pDTI->CreateDTI(sw, sh, sd);

	for (uint z = 0; z < sd; ++z) {
		for (uint y = 0; y < sh; ++y) {
			for (uint x = 0; x < sw; ++x) {
				float fx = (float)x / (float)scale;
				float fy = (float)y / (float)scale;
				float fz = (float)z / (float)scale;
				float dti[7];
				pTemp->GetDTI(fx, fy, fz, dti);
				pDTI->SetDTI(x, y, z, dti);
			}
		}
	}

	SafeDelete(pTemp);
}

float FindNextVoxel(CPoint3F &pos, CVector3F &dir, CPoint3I &voxel)
{
	float u_x = voxel.x + 0.5f;
	float u_y = voxel.y + 0.5f;
	float u_z = voxel.z + 0.5f;
	float l_x = voxel.x - 0.5f;
	float l_y = voxel.y - 0.5f;
	float l_z = voxel.z - 0.5f;

	float min_l = 1000.0f;
	float l = 0.0f;
	CPoint3I next;

	// case 1
	l = (u_x - pos.m_x) / dir.m_x;
	if (l > 0.0f && l < min_l) {
		min_l = l;
		next.x = voxel.x + 1;
		next.y = voxel.y;
		next.z = voxel.z;
	}

	// case 2
	l = (l_x - pos.m_x) / dir.m_x;
	if (l > 0.0f && l < min_l) {
		min_l = l;
		next.x = voxel.x - 1;
		next.y = voxel.y;
		next.z = voxel.z;
	}

	// case 3
	l = (u_y - pos.m_y) / dir.m_y;
	if (l > 0.0f && l < min_l) {
		min_l = l;
		next.x = voxel.x;
		next.y = voxel.y + 1;
		next.z = voxel.z;
	}

	// case 4
	l = (l_y - pos.m_y) / dir.m_y;
	if (l > 0.0f && l < min_l) {
		min_l = l;
		next.x = voxel.x;
		next.y = voxel.y - 1;
		next.z = voxel.z;
	}

	// case 5
	l = (u_z - pos.m_z) / dir.m_z;
	if (l > 0.0f && l < min_l) {
		min_l = l;
		next.x = voxel.x;
		next.y = voxel.y;
		next.z = voxel.z + 1;
	}

	// case 6
	l = (l_z - pos.m_z) / dir.m_z;
	if (l > 0.0f && l < min_l) {
		min_l = l;
		next.x = voxel.x;
		next.y = voxel.y;
		next.z = voxel.z - 1;
	}

	min_l += DELTA;

	pos.m_x += dir.m_x * min_l;
	pos.m_y += dir.m_y * min_l;
	pos.m_z += dir.m_z * min_l;

	voxel.x = next.x;
	voxel.y = next.y;
	voxel.z = next.z;

	return min_l;
}

CVector3F ComputeNextDirection(float *dti, CVector3F in)
{
	CVector3F out;
	float cl = pDTI->DTI2CL(dti);
	CVector3F v;
	pDTI->DTI2Direction(dti, &v);

	if (InnerProduct(in, v) < 0.0f) {
		v.m_x = -v.m_x;
		v.m_y = -v.m_y;
		v.m_z = -v.m_z;
	}

	out.m_x = dti[1] * in.m_x + dti[2] * in.m_y + dti[3] * in.m_z;
	out.m_y = dti[2] * in.m_x + dti[4] * in.m_y + dti[5] * in.m_z;
	out.m_z = dti[3] * in.m_x + dti[5] * in.m_y + dti[6] * in.m_z;
	out.Normalize();

	CVector3F res;

	res.m_x = cl * v.m_x + (1.0f - cl) * (0.8f * in.m_x + 0.2f * out.m_x);
	res.m_y = cl * v.m_y + (1.0f - cl) * (0.8f * in.m_y + 0.2f * out.m_y);
	res.m_z = cl * v.m_z + (1.0f - cl) * (0.8f * in.m_z + 0.2f * out.m_z);

	res.Normalize();

	return res;
}

bool CheckVoxel(CPoint3I &voxel)
{
	int w = pDTI->GetWidth();
	int h = pDTI->GetHeight();
	int d = pDTI->GetDepth();

	if (voxel.x >= 0 && voxel.x < w && 
		voxel.y >= 0 && voxel.y < h && 
		voxel.z >= 0 && voxel.z < d)
		return true;
	else
		return false;
}

void TensorDeflection(Seed &seed, CFiber *fiber, const int x, const int y, const int z)
{
	int fc = 0;
	int bc = 0;
	CPoint3F *pF = new CPoint3F[100];
	CPoint3F *pB = new CPoint3F[100];
	memset(pF, 0, sizeof(CPoint3F)*100);
	memset(pB, 0, sizeof(CPoint3F)*100);

	// forward
	float dti[7];
	pDTI->GetDTI(seed.pos, dti);
	float fa = pDTI->DTI2FA(dti);
	CPoint3F pos = seed.pos;
	CVector3F dir = seed.dir;
	float length = 0.0f;
	CPoint3I voxel = CPoint3I(x, y, z);
	while(fa > faThreshold && length < maxLength) {
		pF[fc] = pos;
		
		length += FindNextVoxel(pos, dir, voxel);
		if(CheckVoxel(voxel) == false)
			break;

		pDTI->GetDTIFast(voxel.x, voxel.y, voxel.z, dti);
		CVector3F next_dir = ComputeNextDirection(dti, dir);
		if (InnerProduct(next_dir, dir) < angleThreshold)
			break;

		dir = next_dir;
		fa = pDTI->DTI2FA(dti);
		++fc;
	}

	// backward
	pDTI->GetDTI(seed.pos, dti);
	fa = pDTI->DTI2FA(dti);
	pos = seed.pos;
	dir = seed.dir;
	dir.m_x = -dir.m_x;	dir.m_y = -dir.m_y;	dir.m_z = -dir.m_z;
	length = 0.0f;
	voxel = CPoint3I(x, y, z);
	while(fa > faThreshold && length < maxLength) {
		pB[bc] = pos;

		length += FindNextVoxel(pos, dir, voxel);
		if(CheckVoxel(voxel) == false)
			break;

		pDTI->GetDTIFast(voxel.x, voxel.y, voxel.z, dti);
		CVector3F next_dir  = ComputeNextDirection(dti, dir);
		if (InnerProduct(next_dir, dir) < angleThreshold)
			break;

		dir = next_dir;
		fa = pDTI->DTI2FA(dti);
		++bc;
	}

	// resample the fiber
	float stepLength = innerSteps * stepSize;
	fiber->m_pF[0] = seed.pos;
	fiber->m_fCount = 1;
	float currentLength = 0.0f;
	for (int i = 1; i < fc; ++i) {
		CVector3F dis = pF[i] - pF[i-1];
		currentLength += dis.Length();
		if (currentLength > stepLength) {
			float delta = currentLength - stepLength;
			float ratio = delta / dis.Length();
			fiber->m_pF[fiber->m_fCount] = pF[i] * (1.0f - ratio) + pF[i-1] * ratio;
			fiber->m_fCount++;
			currentLength -= stepLength;
		}
	}

	fiber->m_pB[0] = seed.pos;
	fiber->m_bCount = 1;
	currentLength = 0.0f;
	for (int i = 1; i < bc; ++i) {
		CVector3F dis = pB[i] - pB[i-1];
		currentLength += dis.Length();
		if (currentLength > stepLength) {
			float delta = currentLength - stepLength;
			float ratio = delta / dis.Length();
			fiber->m_pB[fiber->m_bCount] = pB[i] * (1.0f - ratio) + pB[i-1] * ratio;
			fiber->m_bCount++;
			currentLength -= stepLength;
		}
	}

	fiber->m_seed = seed;


	//fiber->CreateFiber(pF, pB, fc, bc, &seed);

	SafeDeleteArray(pF);
	SafeDeleteArray(pB);

}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	// pre-scale dti volume
	ScaleDTI();

	// fiber tracking by tensor deflection
	uint w = pDTI->GetWidth();
	uint h = pDTI->GetHeight();
	uint d = pDTI->GetDepth();

	FILE *fp = fopen(pFiberPathname, "w");
	fprintf(fp, "%d %d %d %d %f %f\n", w, h, d, (int)scale, stepSize, maxLength);

	std::string pathname = pFiberPathname;
	size_t pos = pathname.find_last_of(SLASH);
	std::string path = pathname.substr(0, pos+1);

	for (uint z = 0; z < d; ++z) {
		printf("Fiber tracking on slice z = %d ... ", z);

		CFiber *fibers = new CFiber[w*h];
		for (uint y = 0; y < h; ++y) {
			for (uint x = 0; x < w; ++x) {
				Seed seed;
				seed.pos = CPoint3F((float)x, (float)y, (float)z);
				float dti[7];
				pDTI->GetDTIFast(x, y, z, dti);
				pDTI->DTI2Direction(dti, &(seed.dir));

				uint index = y * w + x;
				TensorDeflection(seed, &(fibers[index]), x, y, z);
			}
			float ff = float(y+1) / float(h) * 100.0f;
			printf("\rFiber tracking on slice z = %d ... %.2f%%", z, ff);
		}

		char confPathname[PATHNAME_SIZE];
		char vertPathname[PATHNAME_SIZE];
		sprintf(confPathname, "%sslice_z=%08.04f.conf", path.c_str(), (float)z/(float)scale);
		sprintf(vertPathname, "%sslice_z=%08.04f.vert", path.c_str(), (float)z/(float)scale);
		fprintf(fp, "%s\t%s\n", confPathname, vertPathname);

		SaveFibers(confPathname, vertPathname, fibers, w, h);

		SafeDeleteArray(fibers);

		printf("\rFiber tracking on slice z = %d ... 100.00%% Done.\n", z);
	}

	//for (uint y = 238; y < 239; ++y) {
	//	printf("Fiber tracking on slice y = %d ... ", y);

	//	CFiber *fibers = new CFiber[w*d];
	//	for (uint z = 0; z < d; ++z) {
	//		for (uint x = 0; x < w; ++x) {
	//			Seed seed;
	//			seed.pos = CPoint3F((float)x, (float)y, (float)z);
	//			float dti[7];
	//			pDTI->GetDTIFast(x, y, z, dti);
	//			pDTI->DTI2Direction(dti, &(seed.dir));

	//			uint index = z * w + x;
	//			TensorDeflection(seed, &(fibers[index]), x, y, z);
	//		}
	//		float ff = float(z+1) / float(d) * 100.0f;
	//		printf("\rFiber tracking on slice y = %d ... %.2f%%", y, ff);
	//	}

	//	char confPathname[PATHNAME_SIZE];
	//	char vertPathname[PATHNAME_SIZE];
	//	sprintf(confPathname, "%sslice_y=%08.04f.conf", path.c_str(), (float)y/(float)scale);
	//	sprintf(vertPathname, "%sslice_y=%08.04f.vert", path.c_str(), (float)y/(float)scale);
	//	fprintf(fp, "%s\t%s\n", confPathname, vertPathname);

	//	SaveFibers(confPathname, vertPathname, fibers, w, d);

	//	SafeDeleteArray(fibers);

	//	printf("\rFiber tracking on slice y = %d ... 100.00%% Done.\n", y);
	//}

	SafeDelete(pDTI);

	return 0;
}